#!/opt/local/bin/python

# Copyright Stephen Molloy, 2012
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from parseDSKfile import DSKParser
from scipy import nan, isnan
from scipy import optimize as opt
from py4j.java_gateway import JavaGateway, GatewayClient
from numpy import linspace, diff, argmax
import multiprocessing as mp
import subprocess as sp
import socket
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt

class AJDiskAnalysis(object):
    """Instances of this class may be used for AJDisk analyses of klystron models.
    
    obj = AJDiskAnalysis(DSKParsObj = DSKParser(fname))
    It accepts a DSKParser object as input, and uses this to instantiate the
    appropriate klystron and simulation parameters."""
    def __init__(self, DSKParsObj, port=nan):
        self.DSKParsObj = DSKParsObj
        self.port = port
        # First we start a JVM to serve the AJDisk code in a subprocess of this thread
        if isnan(port):
            cmd = ["java", "se/esss/stephenmolloy/AJDiskEntryPoint"]
            self.serverProc = sp.Popen(cmd)
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            # Repeatedly try to connect to the port running the server, only
            # breaking the loop when the server is fully established
            # This blocks the code until the server is in the right state.
            while True:
                try:
                    s.connect(('127.0.0.1', 25333))
                    s.shutdown(2)
                    print "Server opened: port " + str(port)
                    break
                except:
                    pass
            self.gateway = JavaGateway()
        else:
            # Start a JVM on the right port in a subprocess.  Block program
            # execution until the server is established using the same trick as above
            cmd = ["java", "se/esss/stephenmolloy/AJDiskEntryPoint", str(port)]
            self.serverProc = sp.Popen(cmd)
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            while True:
                try:
                    s.connect(('127.0.0.1', port))
                    s.shutdown(2)
                    print "Server opened: port " + str(port)
                    break
                except:
                    pass
            self.gateway = JavaGateway(gateway_client = GatewayClient(port=port))

        # Initialise a new klystron, and set up the operating params, and the various cavities
        self.klysanal = self.gateway.entry_point.getNewKlystron()
        self.klysanal.deleteAllCavities()
        self.klysanal.setOperatingParameters(
                DSKParsObj.opParam.V0,
                DSKParsObj.opParam.I0,
                DSKParsObj.opParam.F0,
                DSKParsObj.opParam.Pin,
                DSKParsObj.opParam.a,
                DSKParsObj.opParam.b)
        
        for cav in DSKParsObj.cavs:
            self.klysanal.addCavity(
                    cav.RoverQ,
                    cav.M,
                    cav.Qe,
                    cav.Qo,
                    cav.f,
                    cav.z,
                    cav.d,
                    cav.harmonic)

    def __getnewargs__(self):
        return (AJDiskAnalysis.__str__(self),)
            
    def calculateAnalytics(self):
        """A method that links directly to the method of the same name in the
        Java-based klystron object."""
        self.klysanal.calculateAnalytics()

    def simulate(self, disks=nan, rings=nan, steps=nan, maxiter=nan):
        """Performs the simulation.  The various inputs allow an override
        of the simulation parameters extracted from the DSKParser object."""
        if isnan(disks):   disks   = self.DSKParsObj.simobj.disks
        if isnan(rings):   rings   = self.DSKParsObj.simobj.rings
        if isnan(steps):   steps   = self.DSKParsObj.simobj.steps
        if isnan(maxiter): maxiter = self.DSKParsObj.simobj.maxiter
        print "Simulating: disks=%i : rings=%i : steps=%i : maxiter=%i" % (disks, rings, steps, maxiter)
        self.klysanal.simulate(disks, rings, steps, maxiter)

    def setBeamVoltage(self, V):
        """Sets the beam voltage to the required value, and alters the beam
        current so as to be correct for a space-charge limited beam.  I.e.
        a constant perveance klystron."""
        newVolts = V
        perv = self.DSKParsObj.opParam.perv
        newAmps  = perv * (V**1.5)
        self.klysanal.setBeamVoltage(newVolts)
        self.klysanal.setBeamCurrent(newAmps)

    def close(self):
        """This method should be called when the user is finished with the current
        klystron.  If not, this will result in a memory leak that can be very bad
        for large optimisations.  This detaches the klystron object, shutdowns the gateway,
        and kills the server.
        Execution is blocked until the server process is reported to be terminated."""
        self.gateway.detach(self.klysanal)
        self.gateway.shutdown()
        self.serverProc.kill()
        self.serverProc.wait()
        print("Server closed: port " + str(self.port))

def parallelFindMaxInRange(lo, hi, num, filename, ports, beamV=nan, mode=0):
    """This method is used by the algorithm to find the saturated power of a klystron.
    It scans the input power over a range given by inputs to the function, and searches
    for the point that the derivative changes sign.
    It returns a tuple of two tuples.  The first is the input powers that bracket the
    location of the maximum point, and the second contains the two output powers bracketing
    this point."""
    scanrange = linspace(lo, hi, num=num)
    pool = mp.Pool(processes=num)
    counter = 0
    result, Pout = [], []
    for power in scanrange:
        result.append(pool.apply_async(parallelKlysOutputPower, 
            (power, filename, ports[counter], beamV)))
        counter += 1
    for i in result:
        Pout.append(i.get())

    # Depending on the mode input parameter, determine the maximum by finding the point
    # at which the derivative changes sign, or the maximum of the curve.
    if mode==0: ind = findFirstMax(Pout)
    else:       ind = argmax(Pout)

    try: bracket = (scanrange[ind], scanrange[ind+1])
    except IndexError:
        diff = scanrange[ind] - scanrange[ind-1]
        bracket = (scanrange[ind], scanrange[ind]+diff)

    Poutbracket = (Pout[ind], Pout[ind+1])
    return (bracket, Poutbracket)

def parallelFindSatPower(filename, ports, beamV=nan):
    """Find the input and output powers where the klystron saturates by scanning
    the input power over some range, and the interpolating to predict the maximum point."""
    iniPinRange = (30, 110)
    PinRange = linspace(iniPinRange[0], iniPinRange[1], num=24)
    PoutRange = parallelPowerScan(Pin=PinRange, filename=filename,
            ports=ports, beamV=beamV)
    interpfunc = interp1d(PinRange, PoutRange, kind=3, bounds_error=True)

    PinDense = linspace(iniPinRange[0], iniPinRange[1], num=1000)
    PoutDense = interpfunc(PinDense)
    maxInd = PoutDense.argmax()
    Pin  = PinDense[maxInd]
    Pout = PoutDense[maxInd]
    return (Pin, Pout)

def findFirstMax(inarray):
    """Helper function parallelFindMaxInRange()"""
    diffarray = diff(inarray)
    for i in range(len(diffarray)):
        if diffarray[i]<0:
            break
    return i

def findSatPower(AJDiskAnalObj):
    """This finds the saturated power of the klystron by searching for the maximum between
    10 and 100 W."""
    Pin_sat = opt.fminbound(negOP, 10.0, 100.0, args=(AJDiskAnalObj, 1.0), xtol=1e-2)
    #Pin_sat = opt.fmin(negOP, 15.0, args=(AJDiskAnalObj, 1.0), ftol=5e3)
    #AJDiskAnalObj.klysanal.setInputPower(Pin_sat[0])
    AJDiskAnalObj.klysanal.setInputPower(Pin_sat)
    Pout_sat = AJDiskAnalObj.klysanal.getSimulatedOutputPower()
    #return (Pin_sat[0], Pout_sat)
    return (Pin_sat, Pout_sat)

def klysOutputPower(Pin, AJDiskAnalObj):
    """Returns the output power of the klystron for a given input power."""
    AJDiskAnalObj.klysanal.setInputPower(float(Pin))
    #AJDiskAnalObj.simulate(disks=20)
    AJDiskAnalObj.simulate()
    return AJDiskAnalObj.klysanal.getSimulatedOutputPower()

def parallelKlysOutputPower(Pin, fname, port, beamV=nan):
    """Returns the output power of the klystron for a given input power."""
    AJDiskAnalObj = AJDiskAnalysis(DSKParsObj = DSKParser(fname=fname), port=port)
    AJDiskAnalObj.klysanal.setInputPower(float(Pin))
    if not isnan(beamV):
        AJDiskAnalObj.setBeamVoltage(V=beamV)
    AJDiskAnalObj.simulate()
    Pout = AJDiskAnalObj.klysanal.getSimulatedOutputPower()
    AJDiskAnalObj.close()
    return Pout

def parallelKlysOutputPhase(Pin, fname, port, beamV=nan):
    """Returns the output power of the klystron for a given input power."""
    AJDiskAnalObj = AJDiskAnalysis(DSKParsObj = DSKParser(fname=fname), port=port)
    AJDiskAnalObj.klysanal.setInputPower(float(Pin))
    if not isnan(beamV):
        AJDiskAnalObj.setBeamVoltage(V=beamV)
    AJDiskAnalObj.simulate()
    numCavs = AJDiskAnalObj.klysanal.getCavityCount()
    Pout = AJDiskAnalObj.klysanal.getSimulatedCavityVoltagePhase(numCavs-1)
    AJDiskAnalObj.close()
    return Pout

def parallelKlysOutput(Pin, fname, port, beamV=nan):
    """Returns the output power of the klystron for a given input power."""
    AJDiskAnalObj = AJDiskAnalysis(DSKParsObj = DSKParser(fname=fname), port=port)
    AJDiskAnalObj.klysanal.setInputPower(float(Pin))
    if not isnan(beamV):
        AJDiskAnalObj.setBeamVoltage(V=beamV)
    AJDiskAnalObj.simulate()
    numCavs = AJDiskAnalObj.klysanal.getCavityCount()
    PhaseOut = AJDiskAnalObj.klysanal.getSimulatedCavityVoltagePhase(numCavs-1)
    PowerOut = AJDiskAnalObj.klysanal.getSimulatedOutputPower()
    AJDiskAnalObj.close()
    return (PowerOut, PhaseOut)

def negOP(Pin, AJDiskAnalObj, waste):
    """A convenience function that returns the negative of the output power.
    This is useful for functions using minimisation routines to maximise
    the klystron power."""
    return -klysOutputPower(Pin, AJDiskAnalObj)

def findDesiredV(reqP, AJDiskAnalObj):
    """Uses a minimisation method to find the necessary beam voltage to achieved
    a required saturated power.
    It is assumed that the klystron has a constant perveance over the optimisation range."""
    V, fval, ierr, numfunc = opt.fminbound(diffFromReqPower, 1e3, 200e3, args=(reqP, AJDiskAnalObj), xtol=1e1, full_output=True)
    return V

def parallelFindDesiredV(reqP, fname, ports):
    """Uses a minimisation method to find the necessary beam voltage to achieved
    a required saturated power.
    It is assumed that the klystron has a constant perveance over the optimisation range."""
    V, fval, ierr, numfunc = opt.fminbound(diffFromReqPower, 1e3, 200e3, args=(reqP, fname, ports), xtol=1e1, full_output=True)
    fhand = open('numfunc_count.txt', 'a')
    fhand.write(str(numfunc) + '\n')
    fhand.close()
    return V

def setVSatPower(V, filename, ports):
    """Find the saturated power at a given beam voltage.
    Returns both the input and output powers at saturation."""
    satPower = parallelFindSatPower(filename=filename, ports=ports, beamV=V)
    print "Saturated power: In = %f W : Out = %e W" % (satPower[0], satPower[1])
    return satPower

def diffFromReqPower(V, reqP, filename, ports):
    """Returns the squared difference between the saturated power and the
    required power."""
    Pin, Pout = setVSatPower(V, filename, ports)
    return (reqP - Pout)**2

def parallelPowerScan(Pin, filename, ports, beamV=nan):
    Pout, result, AJobj = [], [], []
    counter = 0
    pool = mp.Pool(processes=24)
    for power in Pin:
        result.append(pool.apply_async(parallelKlysOutputPower, 
            (power, filename, ports[counter], beamV)))
        counter += 1
    for i in result:
        Pout.append(i.get())
    return Pout

def parallelPowerScanPhaseOut(Pin, filename, ports, beamV=nan):
    Pout, result, AJobj = [], [], []
    counter = 0
    pool = mp.Pool(processes=24)
    for power in Pin:
        result.append(pool.apply_async(parallelKlysOutputPhase, 
            (power, filename, ports[counter], beamV)))
        counter += 1
    for i in result:
        Pout.append(i.get())
    return Pout

def parallelPowerScanPowerPhase(Pin, filename, ports, beamV=nan):
    PowerOut, PhaseOut, result, AJobj = [], [], [], []
    counter = 0
    pool = mp.Pool(processes=24)
    for power in Pin:
        result.append(pool.apply_async(parallelKlysOutput, 
            (power, filename, ports[counter], beamV)))
        counter += 1
    for i in result:
        out = i.get()
        PowerOut.append(out[0])
        PhaseOut.append(out[1])
    return (PowerOut, PhaseOut)

def powerScan(Pin, AJDiskAnalObj):
    Pout = []
    for power in Pin:
        Pout.append( klysOutputPower(power, AJDiskAnalObj) )
    return Pout

